#include "gtd_common.hpp"
#include "gtd_convolution.hpp"

#define KERNEL_RADIUS 8
#define KERNEL_LENGTH (2 * KERNEL_RADIUS + 1)

NAMESPACE_BEG(gtd)

void computeGaussianKernels(float sigma, std::vector<float> order[2])
{
    // precompute constants
    const float doubleSigma2 = 2 * sigma * sigma;
    order[0].resize(KERNEL_LENGTH);
    order[1].resize(KERNEL_LENGTH);

    // compute 0th and 1st order gaussians and their cummulative sums
    float sum0 = 1.0f;
    float sum1 = 0.0f;
    for(int i = 0; i <= KERNEL_RADIUS; ++i)
    {
        order[0][KERNEL_RADIUS + i] = std::exp(-i * i / doubleSigma2);
        order[0][KERNEL_RADIUS - i] = order[0][KERNEL_RADIUS + i];
        order[1][KERNEL_RADIUS + i] = -i * order[0][KERNEL_RADIUS + i];
        order[1][KERNEL_RADIUS - i] = -order[1][KERNEL_RADIUS + i];
        if(i > 0)
        {
            sum0 += 2 * order[0][KERNEL_RADIUS + i];
            sum1 += i * order[1][KERNEL_RADIUS + i] - i * order[1][KERNEL_RADIUS - i];
        }
    }

    // normalize
    for(uint i = 0; i < order[0].size(); ++i)
    {
        order[0][i] /= sum0;
    }
    for(uint i = 0; i < order[1].size(); ++i)
    {
        order[1][i] /= sum1;
    }
}

__constant__ float kernelData[KERNEL_LENGTH];

static void setupArguments(const std::vector<float>& kernel)
{
    static const float* lastKernel = NULL;

    if(lastKernel != &kernel[0])
    {
        CUDA_ERROR_CHECK(cudaMemcpyToSymbol(kernelData, &kernel[0], KERNEL_LENGTH * sizeof(float)));
        lastKernel = &kernel[0];
    }
}

#define ROWS_BLOCKDIM_X 16
#define ROWS_BLOCKDIM_Y 4
#define ROWS_RESULT_STEPS 8
#define ROWS_HALO_STEPS 1

__global__ void convolveHorizontal
(    
    float* dst,
    const float* src,
    uint width,
    uint height
)
{
    __shared__ float s_Data[ROWS_BLOCKDIM_Y][(ROWS_RESULT_STEPS + 2 * ROWS_HALO_STEPS) * ROWS_BLOCKDIM_X];

    //Offset to the left halo edge
    const int baseX = (blockIdx.x * ROWS_RESULT_STEPS - ROWS_HALO_STEPS) * ROWS_BLOCKDIM_X + threadIdx.x;
    const int baseY = blockIdx.y * ROWS_BLOCKDIM_Y + threadIdx.y;

    src += baseY * width + baseX;
    dst += baseY * width + baseX;

    //Load main data
    #pragma unroll
    for(int i = ROWS_HALO_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i++)
        s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] = src[i * ROWS_BLOCKDIM_X];

    //Load left halo
    #pragma unroll
    for(int i = 0; i < ROWS_HALO_STEPS; i++)
        s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] = (baseX >= -i * ROWS_BLOCKDIM_X ) ? src[i * ROWS_BLOCKDIM_X] : *(src - baseX);

    //Load right halo
    #pragma unroll
    for(int i = ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS + ROWS_HALO_STEPS; i++)
        s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X] = (width - baseX > i * ROWS_BLOCKDIM_X) ? src[i * ROWS_BLOCKDIM_X] : src[width - baseX - 1];

    //Compute and store results
    __syncthreads();
    #pragma unroll
    for(int i = ROWS_HALO_STEPS; i < ROWS_HALO_STEPS + ROWS_RESULT_STEPS; i++)
    {
        float sum = 0;

        #pragma unroll
        for(int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++)
            sum += kernelData[KERNEL_RADIUS - j] * s_Data[threadIdx.y][threadIdx.x + i * ROWS_BLOCKDIM_X + j];

        dst[i * ROWS_BLOCKDIM_X] = sum;
    }
}

void convolveHorizontal(float* dst, const float* src, uint width, uint height, const std::vector<float>& kernel)
{
    assert(ROWS_BLOCKDIM_X * ROWS_HALO_STEPS >= KERNEL_RADIUS);
    assert(width % (ROWS_RESULT_STEPS * ROWS_BLOCKDIM_X) == 0);
    assert(height % ROWS_BLOCKDIM_Y == 0);

    setupArguments(kernel);

    dim3 blocks(width / (ROWS_RESULT_STEPS * ROWS_BLOCKDIM_X), height / ROWS_BLOCKDIM_Y);
    dim3 threads(ROWS_BLOCKDIM_X, ROWS_BLOCKDIM_Y);

    convolveHorizontal<<<blocks, threads>>>(dst, src, width, height);
}

#define COLUMNS_BLOCKDIM_X 16
#define COLUMNS_BLOCKDIM_Y 8
#define COLUMNS_RESULT_STEPS 8
#define COLUMNS_HALO_STEPS 1

__global__ void convolveVertical
(    
    float* dst,
    const float* src,
    uint width,
    uint height
)
{
    __shared__ float s_Data[COLUMNS_BLOCKDIM_X][(COLUMNS_RESULT_STEPS + 2 * COLUMNS_HALO_STEPS) * COLUMNS_BLOCKDIM_Y + 1];

    //Offset to the upper halo edge
    const int baseX = blockIdx.x * COLUMNS_BLOCKDIM_X + threadIdx.x;
    const int baseY = (blockIdx.y * COLUMNS_RESULT_STEPS - COLUMNS_HALO_STEPS) * COLUMNS_BLOCKDIM_Y + threadIdx.y;
    src += baseY * width + baseX;
    dst += baseY * width + baseX;

    //Main data
    #pragma unroll
    for(int i = COLUMNS_HALO_STEPS; i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS; i++)
        s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] = src[i * COLUMNS_BLOCKDIM_Y * width];

    //Upper halo
    #pragma unroll
    for(int i = 0; i < COLUMNS_HALO_STEPS; i++)
        s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y] = (baseY >= -i * COLUMNS_BLOCKDIM_Y) ? src[i * COLUMNS_BLOCKDIM_Y * width] : *(src - width * baseY);

    //Lower halo
    #pragma unroll
    for(int i = COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS; i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS + COLUMNS_HALO_STEPS; i++)
        s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y]= (height - baseY > i * COLUMNS_BLOCKDIM_Y) ? src[i * COLUMNS_BLOCKDIM_Y * width] : src[width * (height - baseY - 1)];

    //Compute and store results
    __syncthreads();
    #pragma unroll
    for(int i = COLUMNS_HALO_STEPS; i < COLUMNS_HALO_STEPS + COLUMNS_RESULT_STEPS; i++){
        float sum = 0;
        #pragma unroll
        for(int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++)
            sum += kernelData[KERNEL_RADIUS - j] * s_Data[threadIdx.x][threadIdx.y + i * COLUMNS_BLOCKDIM_Y + j];

        dst[i * COLUMNS_BLOCKDIM_Y * width] = sum;
    }
}

void convolveVertical(float* dst, const float* src, uint width, uint height, const std::vector<float>& kernel)
{
    assert(COLUMNS_BLOCKDIM_Y * COLUMNS_HALO_STEPS >= KERNEL_RADIUS);
    assert(width % COLUMNS_BLOCKDIM_X == 0);
    assert(height % (COLUMNS_RESULT_STEPS * COLUMNS_BLOCKDIM_Y) == 0);

    setupArguments(kernel);

    dim3 blocks(width / COLUMNS_BLOCKDIM_X, height / (COLUMNS_RESULT_STEPS * COLUMNS_BLOCKDIM_Y));
    dim3 threads(COLUMNS_BLOCKDIM_X, COLUMNS_BLOCKDIM_Y);

    convolveVertical<<<blocks, threads>>>(dst, src, width, height);
}

NAMESPACE_END(gtd)
